#include <cutil_inline.h>
#include <device_launch_parameters.h>
#include <device_functions.h>
#include <iostream>
#include <math.h>
#include <stdio.h>

#define MODE 1

#define BSIZE		1024

#if MODE == 0
#define SMEMSIZE	BSIZE*sizeof(int)
#else
#define SMEMSIZE	BSIZE*(sizeof(int)+sizeof(char))
#endif

#define GSIZE(elemNo) \
	elemNo%BSIZE==0?elemNo/BSIZE:elemNo/BSIZE+1;

#define timer_start(start, stop) \
	cudaEventCreate(&start);\
	cudaEventCreate(&stop);\
	cudaEventRecord(start, 0);
#define timer_stop(start, stop, elapsedTime) \
	cudaEventRecord(stop, 0);\
	cudaEventSynchronize(stop);\
	cudaEventElapsedTime(&elapsedTime, start, stop);\
	cudaEventDestroy(start);\
	cudaEventDestroy(stop);\
	std::cout<<elapsedTime<<"ms"<<std::endl;
#define sysinfo() \
	cudaDeviceProp prop; \
	int cuDeviceCount; \
	cudaGetDeviceCount(&cuDeviceCount); \
	for(int i=0; i<cuDeviceCount; i++){ \
		cudaGetDeviceProperties( &prop, i ); \
		std::cout<<"name:"<<prop.name<<std::endl; \
		std::cout<<"Compute capability:" \
			<<prop.major<<"."<<prop.minor<<std::endl; \
		std::cout<<"Total global mem:" \
			<<prop.totalGlobalMem<<std::endl; \
		std::cout<<"Shared mem per mp:" \
			<<prop.sharedMemPerBlock<<std::endl; \
		std::cout<<"Max threads per block:" \
			<<prop.maxThreadsPerBlock<<std::endl; \
		std::cout<<std::endl; \
	}

#if MODE == 0
__device__ int scan_warp(volatile int* data, const int tid){
	const unsigned int lane = tid&31;
	if(lane>=1) data[tid]=data[tid-1]+data[tid];
	if(lane>=2) data[tid]=data[tid-2]+data[tid];
	if(lane>=4) data[tid]=data[tid-4]+data[tid];
	if(lane>=8) data[tid]=data[tid-8]+data[tid];
	if(lane>=16) data[tid]=data[tid-16]+data[tid];
	return data[tid];
}

__device__ int scan_block(volatile int* data, const int tid){
	const unsigned int lane = tid & 31;
	const unsigned int wid	= tid >> 5; //warp id

	int val = scan_warp(data, tid);
	__syncthreads();
	if(lane==31 && tid & (BSIZE-1) != BSIZE-1) data[wid]=data[tid];
	__syncthreads();
	if(tid & (BSIZE-1) == BSIZE-1) data[wid]=data[tid];
	if(wid==0)	 scan_warp(data, tid);
	__syncthreads();
	if(wid>0)	 val=data[wid-1]+val;
	__syncthreads();
	data[tid] = val; 
	__syncthreads();

	return val;
}

extern __shared__ int sdata[];
__global__ void scan_kernel(int *data, int *bres){
	const unsigned int tid	= threadIdx.x;
	const unsigned int i	= threadIdx.x+blockDim.x*blockIdx.x;
	sdata[tid] = data[i];
	scan_block(sdata, tid);
	data[i] = sdata[tid];
	if(tid==0 && bres != NULL)
		bres[blockIdx.x+1]=sdata[BSIZE-1];
	if(i==0 && bres != NULL)
		bres[0]=0;
}

__global__ void addback_kernel(int* data, int *bres){
	const unsigned int tid = threadIdx.x;
	const unsigned int i	= threadIdx.x+blockDim.x*blockIdx.x;
	sdata[tid]=data[i];
	if(tid==0) sdata[BSIZE]=bres[blockIdx.x];
	__syncthreads();
	sdata[tid]=sdata[tid]+sdata[BSIZE];
	data[i]=sdata[tid];
}

int main(int argc, char** argv){
	float time;
	int *data, *hdata, *bres;
	int gsize = GSIZE(ELEMNO);

	cudaEvent_t start, stop;
	hdata = (int*)malloc(ELEMNO*sizeof(int));

	for(int i=0; i<ELEMNO; i++)
		hdata[i]=1;
	std::cout<<hdata[ELEMNO-1];
	cudaMalloc((void**)&bres, (gsize+1)*sizeof(int));
	cudaMalloc((void**)&data, ELEMNO*sizeof(int));

	timer_start(start, stop);
	for(int i=0; i<100; i++){
		cudaMemcpy(data, hdata, ELEMNO*sizeof(int), cudaMemcpyHostToDevice);
		scan_kernel<<<gsize, BSIZE, SMEMSIZE>>>(data, bres);
		int bresgsize=GSIZE(gsize);
		scan_kernel<<<bresgsize, BSIZE, SMEMSIZE>>>(bres, NULL);
		int smemsize=SMEMSIZE+sizeof(int);
		addback_kernel<<<gsize, BSIZE, smemsize>>>(data, bres);
	}
	timer_stop(start, stop, time);
	cudaMemcpy(hdata, data, ELEMNO*sizeof(int), cudaMemcpyDeviceToHost);
	//cudaMemcpy(hdata, bres, gsize*sizeof(int), cudaMemcpyDeviceToHost);
	std::cout<<hdata[ELEMNO-1]<<std::endl;
}

#else
__device__ int scan_warp(volatile int* data, volatile char* flag, const int tid){
	const unsigned int lane = tid&31;
	if(lane >= 1) if(flag[tid] == flag[tid-1]) data[tid] = data[tid-1] + data[tid];
	if(lane >= 2) if(flag[tid] == flag[tid-2]) data[tid] = data[tid-2] + data[tid];
	if(lane >= 4) if(flag[tid] == flag[tid-4]) data[tid] = data[tid-4] + data[tid];
	if(lane >= 8) if(flag[tid] == flag[tid-8]) data[tid] = data[tid-8] + data[tid];
	if(lane >= 16) if(flag[tid] == flag[tid-16]) data[tid] = data[tid-16] + data[tid];
	return data[tid];
}

__device__ int scan_block(volatile int* data, volatile char* flag, const int tid){
	const unsigned int lane = tid & 31;
	const unsigned int wid	= tid >> 5; //warp id

	int vdata = scan_warp(data, flag, tid);
	int vflag = flag[tid];
	__syncthreads();
	if(lane == 31 && tid & (BSIZE-1) != BSIZE-1) {
		data[wid] = data[tid]; flag[wid] = flag[tid];}
	__syncthreads();
	if(tid & (BSIZE-1) == BSIZE-1) {
		data[wid] = data[tid]; flag[wid] = flag[tid];}
	if(wid == 0)
		scan_warp(data, flag, tid);
	__syncthreads();
	if(wid>0) 
		if(vflag == flag[wid-1]) vdata = data[wid-1]+vdata;
	__syncthreads();
	data[tid] = vdata; 
	__syncthreads();

	return vdata;
}

extern __shared__ int smem[];
__global__ void scan_kernel(int *data, char *flag, int *tdata, char *tflag){
	const unsigned int tid	= threadIdx.x;
	const unsigned int i	= threadIdx.x+blockDim.x*blockIdx.x;
	int *sdata = smem;
	char *sflag = (char*)&(smem[BSIZE]);

	sdata[tid] = data[i];
	sflag[tid] = flag[i];
	scan_block(sdata, sflag, tid);
	//scan_warp(sdata, sflag, tid);
	if(tid == 0 && tdata != NULL){
		tdata[blockIdx.x+1] = sdata[BSIZE-1];
		tflag[blockIdx.x+1] = sflag[BSIZE-1];
	}
	if(i == 0 && tdata != NULL){
		tdata[0] = 0;
		tflag[0] = 0;
	}
	data[i] = sdata[tid];
}

__global__ void addback_kernel(int* data, char *flag, int *tdata, char *tflag){
	const unsigned int tid = threadIdx.x;
	const unsigned int i	= threadIdx.x+blockDim.x*blockIdx.x;
	int *sdata = smem;
	char *sflag = (char*)&smem[BSIZE+1];

	sdata[tid] = data[i];
	sflag[tid] = flag[i];
	if(tid==0) {
		sdata[BSIZE] = tdata[blockIdx.x];
		sflag[BSIZE] = tflag[blockIdx.x];
	}
	__syncthreads();
	if(sflag[tid] == sflag[BSIZE])
		sdata[tid] = sdata[tid]+sdata[BSIZE];
	data[i] = sdata[tid];
}

int main(int argc, char** argv){
	const int ELEMNO = 1024;
	float time;
	int *data, *hdata, *tdata;
	char *flag, *hflag, *tflag;
	int gsize = GSIZE(ELEMNO);
	const int BASESIZE = 10;
	cudaEvent_t start, stop;

	hdata = (int*)malloc(ELEMNO*sizeof(int));
	hflag = (char*)malloc(ELEMNO*sizeof(char));

	for(int i=0; i<ELEMNO; i++){
		hdata[i]=1;
		hflag[i]=i/BASESIZE;
	}
	std::cout<<hdata[ELEMNO-1]<<std::endl;

	cudaMalloc((void**)&data, ELEMNO*sizeof(int));
	cudaMalloc((void**)&flag, ELEMNO*sizeof(char));
	cudaMalloc((void**)&tdata, gsize*sizeof(int));
	cudaMalloc((void**)&tflag, gsize*sizeof(char));
	cudaMemcpy(data, hdata, ELEMNO*sizeof(int), cudaMemcpyHostToDevice);
	cudaMemcpy(flag, hflag, ELEMNO*sizeof(char), cudaMemcpyHostToDevice);

	//scan_kernel<<<gsize, BSIZE, SMEMSIZE>>>(data, flag, NULL, NULL);
	scan_kernel<<<gsize, BSIZE, SMEMSIZE>>>(data, flag, tdata, tflag);
	int tempgsize = GSIZE(gsize);
	scan_kernel<<<tempgsize, BSIZE, SMEMSIZE>>>(tdata, tflag, NULL, NULL);
	int smemsize = SMEMSIZE+sizeof(int)+sizeof(char);
	addback_kernel<<<gsize, BSIZE, smemsize>>>(data, flag, tdata, tflag);

	cudaMemcpy(hdata, data, BSIZE*sizeof(int), cudaMemcpyDeviceToHost);
	for(int i=0; i<ELEMNO; i++){
		std::cout<<hdata[i]<<" ";
		if(hdata[i]==BASESIZE)
			std::cout<<std::endl;
	}
}
#endif